from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import StructType, StructField, IntegerType, StringType

if __name__ == '__main__':
    spark = SparkSession.builder.appName("udf_learn").master("local[*]").getOrCreate()
    schema = StructType(
        [
            StructField("id", IntegerType(), nullable=True),
            StructField("name", StringType(), nullable=True),
            StructField("age", IntegerType(), nullable=True)
        ]
    )
    df = spark.read.option("multiline", "true").schema(schema).json("../data/user.json")
    df.printSchema()

    def age_scale(age: int):
        if age <= 20:
            return '0~20'
        elif age <= 40:
            return '21~40'
        else:
            return '40+'


    # 只能用于dsl中使用
    age_scale_udf = udf(age_scale, StringType())
    df.createOrReplaceTempView("user")
    # 用此方式注册udf函数，dsl和sql中都能使用。注意：注册名（sql中用）和返回值（dsl中用）名要一致
    age_scale_udf2 = spark.udf.register("age_scale_udf2", age_scale)

    df.select("id", "name", "age", age_scale_udf("age").alias("age_scale"),
              age_scale_udf2("age").alias("age_scale2")).show()
    spark.sql("select id, name, age, age_scale_udf2(age) as age_scale2 from user").show()
    spark.stop()
